import os
import pandas as pd
import chardet

"""
 * 按照文件大小分割大文件
 * @param input_file 输入文件路径
 * @param output_dir 输出目录
 * @param chunk_size 分割块大小（字节）
 ***/
"""


def split_file(input_file_path, output_dir_path, chunk_size=50 * 1024 * 1024):  # 默认50MB
    """
    按指定大小流式分割文件
    :param input_file_path: 输入文件路径
    :param output_dir_path: 输出目录
    :param chunk_size: 分割块大小（字节）
    """
    if not os.path.exists(output_dir_path):
        os.makedirs(output_dir_path)

    file_number = 0
    with open(input_file_path, 'rb') as f:
        chunk = f.read(chunk_size)
        while chunk:
            output_path = os.path.join(output_dir_path, f"{os.path.basename(input_file_path)}.part{file_number}")
            with open(output_path, 'wb') as out_file:
                out_file.write(chunk)
            file_number += 1
            chunk = f.read(chunk_size)

    print(f"共分割成 {file_number} 个文件，保存在 '{output_dir_path}'")


"""
 * 按行数分割csv文件
 * @param input_file 输入文件路径
 * @param output_dir 输出目录
 * @param rows_per_file 每个文件包含的行数
"""


def split_csv_by_rows(input_file_path, output_dir_path, rows_per_file=10000):
    """
    流式按行数分割CSV文件
    :param input_file_path: 输入CSV文件路径
    :param output_dir_path: 输出目录
    :param rows_per_file: 每个文件包含的行数
    """
    if not os.path.exists(output_dir_path):
        os.makedirs(output_dir_path)

    # 1. 检测原始编码
    with open(input_file_path, 'rb') as f:
        result = chardet.detect(f.read(10000))
    source_encoding = result['encoding']
    print(f"源文件的编码方式为：{source_encoding}")

    if source_encoding is None:
        raise ValueError("无法识别文件编码")

    file_counter = 1
    # 使用迭代器方式逐块读取
    for chunk in pd.read_csv(
            input_file_path,
            encoding=source_encoding,
            chunksize=rows_per_file,
            low_memory=False,
            on_bad_lines='skip'):
        base_name = os.path.splitext(os.path.basename(input_file_path))[0]  # 去掉 .csv
        new_filename = f"{base_name}-part-{file_counter}.csv"
        output_path = os.path.join(output_dir_path, new_filename)
        chunk.to_csv(output_path, index=False, encoding=source_encoding)
        print(f"已生成第 {file_counter} 个文件，保存在 '{output_dir_path}'")
        file_counter += 1

    print(f"共分割成 {file_counter - 1} 个文件，保存在 '{output_dir_path}'")


if __name__ == '__main__':
    input_file = 'D:\\2-work\\25-创意组\\标签用户体系\\标签体系原始数据\\segmentation\\船视宝2025年用户接口访问记录\\船视宝2025年用户接口访问记录-part-4.csv'
    output_dir = 'D:\\2-work\\25-创意组\\标签用户体系\\标签体系原始数据\\segmentation\\船视宝2025年用户接口访问记录\\part4'
    split_csv_by_rows(input_file, output_dir, 200000)
